package com.zelix.tool;

import com.arthas.client.ArthasClient;
import com.arthas.client.ArthasTool;
import com.zelix.bean.AnalyClass;
import com.zelix.bean.AnalyMethod;
import jdk.internal.org.objectweb.asm.commons.Method;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.ClassWriter;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.tree.*;

import java.io.*;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.jar.JarFile;
import java.util.zip.ZipEntry;

public class AnalyTool {
    /**
     * key文件初始化key数据
     *
     * @param keyPath
     * @return
     */
    public static HashSet<String> getKeys(String keyPath) {
        HashSet<String> strs = new HashSet<>();
        try {
            FileInputStream fis = new FileInputStream(keyPath);
            BufferedReader br = new BufferedReader(new InputStreamReader(fis));
            String tmp;
            while ((tmp = br.readLine()) != null) {
                strs.add(tmp);
            }
            br.close();
            fis.close();
        } catch (FileNotFoundException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return strs;
    }

    /**
     * 获取指定类拦截到的所有的方法
     *
     * @param cls
     * @param keys
     * @return
     */
    public static AnalyClass getAnalyClass(String cls, HashSet<String> keys) {
        AnalyClass analyClass = new AnalyClass(cls);
        for (String key : keys) {
            String[] args = key.split(":");
            if (args.length == 4 && cls.equals(args[0])) {
                analyClass.putMethod(args[1] + args[2], new AnalyMethod(key));
            }
        }
        return analyClass;
    }

    /**
     * 获取所有类及其方法
     *
     * @param keys
     * @return
     */
    public static HashMap<String, AnalyClass> getAnalyClass(HashSet<String> keys) {
        HashMap<String, AnalyClass> classes = new HashMap<>();
        AnalyClass analyClass;
        for (String key : keys) {
            String[] args = key.split(":");
            if (args.length == 4) {
                if (classes.get(args[0]) == null) {
                    analyClass = new AnalyClass(args[0]);
                    analyClass.putMethod(args[1] + args[2], new AnalyMethod(key));
                    classes.put(args[0], analyClass);
                } else {
                    analyClass = classes.get(args[0]);
                    analyClass.putMethod(args[1] + args[2], new AnalyMethod(key));
                }
            }
        }

        return classes;
    }

    /**
     * 移除类似"@String[REF_GET_FIELD]"标签
     *
     * @param cmd 返回的命令
     * @return
     */
    public static String removeStringLebal(String cmd) {
        return cmd.substring("@String[".length(), cmd.length() - 1);
    }

    public static void convert(ArthasClient client, JarFile jarFile, AnalyClass analyClass) throws IOException {
        //获取类所在的classloader,idea中的一些方法不加classloader去运行,可能找不到类
        String classLoader = ArthasTool.getClassLoader(client, analyClass.getCls());
        if (classLoader == null) {
            System.out.println("类 " + analyClass.getCls() + " 还未加载,请点击相应功能触发其加载.");
            return;
        }
        //字符串加密的ognl表示方式,一个类种一致,都是这个方法解密字符串
        String stringMethod = null;
        String classPath = analyClass.getCls() + ".class";
        ZipEntry entry = jarFile.getEntry(classPath);


        if (entry != null) {
            ClassReader cr = new ClassReader(jarFile.getInputStream(entry));
            ClassNode cn = new ClassNode();
            cr.accept(cn, 0);

            HashMap<String, AnalyMethod> analyMethods = analyClass.getMethods();
            List<MethodNode> methods = cn.methods;
            for (MethodNode methodNode : methods) {
                if (analyMethods.get(methodNode.name + methodNode.desc) != null) {
                    AnalyMethod method = analyMethods.get(methodNode.name + methodNode.desc);

                    //原始指令
                    InsnList srcList = methodNode.instructions;
                    //新指令,不动原始指令,以免影响list变化
                    InsnList dstList = new InsnList();
                    //遍历指令
                    for (int i = 0; i < srcList.size(); i++) {

                        //存入新的指令中
                        dstList.add(srcList.get(i));

                        //查找invokedynamic指令
                        if (srcList.get(i) instanceof InvokeDynamicInsnNode) {
                            InvokeDynamicInsnNode dynamicNode = (InvokeDynamicInsnNode) srcList.get(i);

                            //查找是不是调用本类内的方法,是的话一般是字符串加密用途
                            if (dynamicNode.bsm.getOwner().equals(method.getCls())) {

                                //首次遇到该命令,初始化stringMethod,以便下面ognl命令使用
                                if (stringMethod == null) {
                                    //遍历本类方法
                                    for (MethodNode mtd : methods) {
                                        //查找符合动态方法描述的方法
                                        if (mtd.desc.equals(dynamicNode.desc)) {
                                            if (mtd.desc.equals("(IJ)Ljava/lang/String;")) {
                                                //这样处理方式是为了处理其他方法时不必修改固定的值了,代码不用修改可以复用
                                                stringMethod = "@" + method.getCls().replaceAll("/", ".") + "@" + mtd.name;
                                            } else {
                                                //其他情况输出再分析吧
                                                System.out.println("非字符串加密用途:" + mtd.name + ":" + mtd.desc);
                                            }
                                        }
                                    }
                                }

                                //stringMethod不为null,开始处理
                                //再加个判断吧,或许真有其他类型方法,增强程序健壮性
                                if (dynamicNode.desc.equals("(IJ)Ljava/lang/String;")) {
                                    /**
                                     * SIPUSH 27134
                                     * LDC 7612129396686384391L
                                     * LLOAD 1
                                     * LXOR
                                     * INVOKEDYNAMIC r (IJ)Ljava/lang/String; handle[H_INVOKESTATIC com/zelix/ZKM.b(Ljava/lang/invoke/MethodHandles$Lookup;Ljava/lang/String;Ljava/lang/invoke/MethodType;)Ljava/lang/invoke/CallSite;] args[]
                                     * 完全匹配上述模式,如果不符,暂不处理.可能存在该加密方法,但是被标签分割,不修改,修改可能导致代码错乱
                                     */
                                    if (srcList.get(i - 1) instanceof InsnNode && srcList.get(i - 1).getOpcode() == Opcodes.LXOR) {//判断上一个是不是LOR
                                        if (srcList.get(i - 2) instanceof VarInsnNode && srcList.get(i - 2).getOpcode() == Opcodes.LLOAD) {
                                            if (srcList.get(i - 3) instanceof LdcInsnNode) {
                                                if (srcList.get(i - 4) instanceof IntInsnNode) {
                                                    LdcInsnNode ldcNode = (LdcInsnNode) srcList.get(i - 3);
                                                    IntInsnNode intNode = (IntInsnNode) srcList.get(i - 4);

                                                    String rtn = client.sendCommand("ognl '" + stringMethod + "(" + intNode.operand + "," + ldcNode.cst + "L^" + method.getKey() + "L)' -c " + classLoader);
                                                    rtn = removeStringLebal(rtn);

                                                    //替换指令
                                                    InsnList tmpInsnList = new InsnList();
                                                    tmpInsnList.add(new LdcInsnNode(rtn));

                                                    //注意！！！！先插入再删除,不然代码会混乱,暂时不明白原因
                                                    //插入新的指令替换
                                                    dstList.add(tmpInsnList);

                                                    //新指令中删除这几条指令
                                                    dstList.remove(srcList.get(i));
                                                    dstList.remove(srcList.get(i - 1));
                                                    dstList.remove(srcList.get(i - 2));
                                                    dstList.remove(srcList.get(i - 3));
                                                    dstList.remove(srcList.get(i - 4));
                                                } else {
                                                    //基本为SIPUSH,也可能为BIPUSH，或者ICONST_0等等,输出再分析把
                                                    System.out.println("其他指令:" + srcList.get(i - 4));
                                                }
                                            }
                                        }
                                    }
                                }

                            } else {
                                //非本类方法,一般为调用其他方法或者类的get和set设置属性
                                //这里固定吧，因为需要编写自定义方法解析参数
                                if (dynamicNode.bsm.getOwner().equals("com/zelix/aae") && dynamicNode.bsm.getName().equals("a") && dynamicNode.bsm.getDesc().equals("(Ljava/lang/invoke/MethodHandles$Lookup;Ljava/lang/String;Ljava/lang/invoke/MethodType;)Ljava/lang/invoke/CallSite;")) {
                                    /**
                                     * LDC -6775250295146407232L
                                     * LLOAD 1
                                     * INVOKEDYNAMIC h (Ljava/lang/StringBuilder;JJJ)Ljava/lang/StringBuilder; handle[H_INVOKESTATIC com/zelix/aae.a(Ljava/lang/invoke/MethodHandles$Lookup;Ljava/lang/String;Ljava/lang/invoke/MethodType;)Ljava/lang/invoke/CallSite;] args[]
                                     * 完全匹配上述模式,如果不符,暂不处理.
                                     */
                                    if (srcList.get(i - 1) instanceof VarInsnNode && srcList.get(i - 1).getOpcode() == Opcodes.LLOAD) {
                                        if (srcList.get(i - 2) instanceof LdcInsnNode) {
                                            LdcInsnNode ldcNode = (LdcInsnNode) srcList.get(i - 2);

                                            String rtn = client.sendCommand("ognl '@com.zelix.tool.Tool@a(\"" + dynamicNode.name + "\"," + ldcNode.cst + "L," + method.getKey() + "L)' -c " + classLoader);
                                            rtn = removeStringLebal(rtn);

                                            //替换指令
                                            InsnList tmpInsnList = new InsnList();
                                            //分析
                                            if (rtn.indexOf('(') == -1) {//不是方法是属性
                                                String[] strs = getOwnerAndName(rtn).split("\\|");
                                                String owner = strs[0];
                                                String name = strs[1];
                                                //从dynamicNode获取属性描述更准确
                                                String desc = dynamicNode.desc.replace("JJ)", ")");
                                                if (rtn.startsWith("findGetter")) {
                                                    //先去掉调用对象，再获取括号内参数
                                                    if (desc.startsWith("(L")) {//第一个参数为对象
                                                        //去掉第一个对象
                                                        desc = "(" + desc.substring(desc.indexOf(";") + 1);
                                                    }
                                                    desc = desc.replace("()", "");
                                                    tmpInsnList.add(new FieldInsnNode(Opcodes.GETFIELD, owner, name, desc.replace("()", "")));
                                                } else if (rtn.startsWith("findSetter")) {
                                                    //先去掉调用对象，再获取括号内参数
                                                    if (desc.startsWith("(L")) {//第一个参数为对象
                                                        //去掉第一个对象
                                                        desc = "(" + desc.substring(desc.indexOf(";") + 1);
                                                    }
                                                    desc = desc.substring(1, desc.indexOf(')'));
                                                    tmpInsnList.add(new FieldInsnNode(Opcodes.PUTFIELD, owner, name, desc.replace("()", "")));
                                                } else if (rtn.startsWith("findStaticGetter")) {
                                                    //getter方法,没有参数,只有返回类型,需要去掉括号
                                                    desc = desc.replace("()", "");
                                                    tmpInsnList.add(new FieldInsnNode(Opcodes.GETSTATIC, owner, name, desc.replace("()", "")));
                                                } else if (rtn.startsWith("findStaticSetter")) {
                                                    //属性设置描述有括号和参数,返回类型是V,需要获取括号内参数类型
                                                    desc = desc.substring(1, desc.indexOf(')'));
                                                    tmpInsnList.add(new FieldInsnNode(Opcodes.PUTSTATIC, owner, name, desc));
                                                }
                                            } else {//方法
                                                String[] strs = analyMethod(rtn).split("\\|");
                                                String owner = strs[0];
                                                String name = strs[1];
                                                //如果有问题可以考虑从dynamicNode获取方法描述,因为asm的Method解析没有看到对繁泛型的处理
                                                //弃用asm的Method解析
                                                //String desc = strs[2];
                                                //去掉invokedynamic指令的JJ参数
                                                String desc = dynamicNode.desc.replace("JJ)", ")");

                                                //invokedynamic里调用的第一个参数,static里没有,非静态方法调用必须要有一个对象,指令是对象描述,rtn返回的owner为方法描述
                                                if (rtn.startsWith("findStatic")) {
                                                    //无需修改,static静态方法不需要对象调用
                                                    tmpInsnList.add(new MethodInsnNode(Opcodes.INVOKESTATIC, owner, name, desc, false));
                                                } else if (rtn.startsWith("findVirtual")) {
                                                    //非static方法不需要对象调用
                                                    if (desc.startsWith("(L")) {//第一个参数为对象
//                                                  //去掉第一个对象
                                                        desc = "(" + desc.substring(desc.indexOf(";") + 1);
                                                    }
                                                    tmpInsnList.add(new MethodInsnNode(Opcodes.INVOKEVIRTUAL, owner, name, desc, false));
                                                } else if (rtn.startsWith("findSpecial")) {
                                                    if (desc.startsWith("(L")) {//第一个参数为对象
//                                                  //去掉第一个对象
                                                        desc = "(" + desc.substring(desc.indexOf(";") + 1);
                                                    }
                                                    tmpInsnList.add(new MethodInsnNode(Opcodes.INVOKEVIRTUAL, owner, name, desc, false));
                                                }
                                            }

                                            dstList.add(tmpInsnList);

                                            dstList.remove(srcList.get(i));
                                            dstList.remove(srcList.get(i - 1));
                                            dstList.remove(srcList.get(i - 2));
                                        }
                                    }
                                } else {
                                    System.out.println(dynamicNode.bsm.getOwner());
                                    System.out.println(dynamicNode.bsm.getName());
                                    System.out.println(dynamicNode.bsm.getDesc());
                                }
                            }
                        }
                    }

                    //新指令替换
                    srcList.clear();
                    srcList.add(dstList);
                    //结束对匹配到的方法的修改
                    methodNode.visitEnd();
                }
            }

            //写入文件
            ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_MAXS);
            cn.accept(cw);
            AnalyTool.classToFile(classPath.substring(classPath.lastIndexOf('/') + 1), cw.toByteArray());
        }
    }

    /**
     * 将class二进制数据写入文件
     *
     * @param name
     * @param data
     */
    public static void classToFile(String name, byte[] data) {
        System.out.println("开始写入文件:" + name);
        File file = new File("./classes/" + name);
        if (!file.getParentFile().exists()) {
            new File("classes").mkdirs();
        }
        if (!file.exists()) {
            try {
                file.createNewFile();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        try {
            FileOutputStream fos = new FileOutputStream(file);
            fos.write(data);
            fos.close();
        } catch (FileNotFoundException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 解析方法
     *
     * @param client
     * @param jarFile
     * @param method
     * @throws IOException
     */
    public static void analy(ArthasClient client, JarFile jarFile, AnalyMethod method) throws IOException {
        //获取类所在的classloader,idea中的一些方法不加classloader去运行,可能找不到类
        String classLoader = ArthasTool.getClassLoader(client, method.getCls());
        if (classLoader == null) {
            System.out.println("类 " + method.getCls() + " 还未加载,请点击相应功能触发其加载.");
            return;
        }
        //字符串加密的ognl表示方式
        String stringMethod = null;
        ZipEntry entry = jarFile.getEntry(method.getCls() + ".class");

        if (entry != null) {
            ClassReader cr = new ClassReader(jarFile.getInputStream(entry));
            ClassNode cn = new ClassNode();
            cr.accept(cn, 0);

            List<MethodNode> methods = cn.methods;
            for (MethodNode methodNode : methods) {
                if (methodNode.name.equals(method.getMethodName()) && methodNode.desc.equals(method.getMethodDesc())) {//查找需要解析的方法
                    InsnList insnList = methodNode.instructions;
                    //遍历指令
                    for (int i = 0; i < insnList.size(); i++) {
                        //查找invokedynamic指令
                        if (insnList.get(i) instanceof InvokeDynamicInsnNode) {
                            InvokeDynamicInsnNode dynamicNode = (InvokeDynamicInsnNode) insnList.get(i);
                            //查找是不是调用本类内的方法,是的话一般是字符串加密用途
                            if (dynamicNode.bsm.getOwner().equals(method.getCls())) {
                                //首次遇到该命令,初始化stringMethod,以便下面ognl命令使用
                                if (stringMethod == null) {
                                    //遍历本类方法
                                    for (MethodNode mtd : methods) {
                                        //查找符合动态方法描述的方法
                                        if (mtd.desc.equals(dynamicNode.desc)) {
                                            if (mtd.desc.equals("(IJ)Ljava/lang/String;")) {
                                                //这样处理方式是为了处理其他方法时不必修改固定的值了,代码不用修改可以复用
                                                stringMethod = "@" + method.getCls().replaceAll("/", ".") + "@" + mtd.name;
                                            } else {
                                                //其他情况输出再分析吧
                                                System.out.println("非字符串加密用途:" + mtd.name + ":" + mtd.desc);
                                            }
                                        }
                                    }
                                }

                                //stringMethod不为null,开始处理
                                //再加个判断吧,或许真有其他类型方法,增强程序健壮性
                                if (dynamicNode.desc.equals("(IJ)Ljava/lang/String;")) {
                                    /**
                                     * SIPUSH 27134
                                     * LDC 7612129396686384391L
                                     * LLOAD 1
                                     * LXOR
                                     * INVOKEDYNAMIC r (IJ)Ljava/lang/String; handle[H_INVOKESTATIC com/zelix/ZKM.b(Ljava/lang/invoke/MethodHandles$Lookup;Ljava/lang/String;Ljava/lang/invoke/MethodType;)Ljava/lang/invoke/CallSite;] args[]
                                     * 完全匹配上述模式,如果不符,暂不处理.可能存在该加密方法,但是被标签分割,不修改,修改可能导致代码错乱
                                     */
                                    if (insnList.get(i - 1) instanceof InsnNode && insnList.get(i - 1).getOpcode() == Opcodes.LXOR) {//判断上一个是不是LOR
                                        if (insnList.get(i - 2) instanceof VarInsnNode && insnList.get(i - 2).getOpcode() == Opcodes.LLOAD) {
                                            if (insnList.get(i - 3) instanceof LdcInsnNode) {
                                                if (insnList.get(i - 4) instanceof IntInsnNode) {
                                                    LdcInsnNode ldcNode = (LdcInsnNode) insnList.get(i - 3);
                                                    IntInsnNode intNode = (IntInsnNode) insnList.get(i - 4);

                                                    String rtn = client.sendCommand("ognl '" + stringMethod + "(" + intNode.operand + "," + ldcNode.cst + "L^" + method.getKey() + "L)' -c " + classLoader);
                                                    System.out.println(rtn);
                                                } else {
                                                    //基本为SIPUSH,也可能为BIPUSH，或者ICONST_0等等,输出再分析把
                                                    System.out.println("其他指令:" + insnList.get(i - 4));
                                                }
                                            }
                                        }
                                    }
                                }
                            } else {
                                //非本类方法,一般为调用其他方法或者类的get和set设置属性
                                //这里固定吧，因为需要编写自定义方法解析参数
                                if (dynamicNode.bsm.getOwner().equals("com/zelix/aae") && dynamicNode.bsm.getName().equals("a") && dynamicNode.bsm.getDesc().equals("(Ljava/lang/invoke/MethodHandles$Lookup;Ljava/lang/String;Ljava/lang/invoke/MethodType;)Ljava/lang/invoke/CallSite;")) {
                                    /**
                                     * LDC -6775250295146407232L
                                     * LLOAD 1
                                     * INVOKEDYNAMIC h (Ljava/lang/StringBuilder;JJJ)Ljava/lang/StringBuilder; handle[H_INVOKESTATIC com/zelix/aae.a(Ljava/lang/invoke/MethodHandles$Lookup;Ljava/lang/String;Ljava/lang/invoke/MethodType;)Ljava/lang/invoke/CallSite;] args[]
                                     * 完全匹配上述模式,如果不符,暂不处理.
                                     */
                                    if (insnList.get(i - 1) instanceof VarInsnNode && insnList.get(i - 1).getOpcode() == Opcodes.LLOAD) {
                                        if (insnList.get(i - 2) instanceof LdcInsnNode) {
                                            LdcInsnNode ldcNode = (LdcInsnNode) insnList.get(i - 2);

                                            String rtn = client.sendCommand("ognl '@com.zelix.tool.Tool@a(\"" + dynamicNode.name + "\"," + ldcNode.cst + "L," + method.getKey() + "L)' -c " + classLoader);
                                            System.out.println(rtn);
                                        }
                                    }
                                } else {
                                    System.out.println(dynamicNode.bsm.getOwner());
                                    System.out.println(dynamicNode.bsm.getName());
                                    System.out.println(dynamicNode.bsm.getDesc());
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    /**
     * 获取方法或者属性的owner和name
     *
     * @param str
     * @return
     */
    public static String getOwnerAndName(String str) {
        String[] strs = str.split(" ");
        //获取最后一个
        str = strs[strs.length - 1];
        //取括号,函数参数是类的话，也会有点，所以先去点参数内容，以免参数的点影响获取结果
        int index = str.lastIndexOf('(');
        if (index != -1) {//存在括号说明是函数，则去点参数内容
            str = str.substring(0, index);
        }
        index = str.lastIndexOf('.');
        //获取name
        String name = str.substring(index + 1, str.length());
        //方法或者属性前所属类
        String owner = str.substring(0, index);
        //替换未字节码识别的斜杠式
        owner = owner.replaceAll("\\.", "/");
        return owner + "|" + name;
    }


    /**
     * 解析java代码格式的方法定义,如public void java.security.cert.PKIXParameters.setDate(java.util.Date)
     * 返回owner|namr|desc描述方式
     *
     * @param str
     * @return
     */
    public static String analyMethod(String str) {
        //获取方法前内容,因为有的方法会抛出异常,这里去掉异常信息
        str = str.substring(0, str.indexOf(')') + 1);
        String[] strs = str.split(" ");
        //暂不判断数组长度,长度肯定不能小于2
        //最后一个数组为方法描述   strs[strs.length - 1] 如java.security.cert.PKIXParameters.setDate(java.util.Date)
        //倒数第二个为方法返回    strs[strs.length - 2]   如void
        //返回NameAndOwner
        //System.out.println(getOwnerAndName(strs[strs.length - 1]));
        //用asm自带方法,只支持返回类型+方法名+参数类型,前面不能有public等修饰,不然识别为返回类型
        Method method = Method.getMethod(strs[strs.length - 2] + " " + strs[strs.length - 1]);
        return getOwnerAndName(strs[strs.length - 1]) + "|" + method.getDescriptor();
    }
}
